{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "epochs = 15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 2 - MultiLayer Split Neural Network\n",
    "\n",
    "<b>Recap:</b> The previous tutorial looked at building a basic SplitNN, where an NN was split into two segments on two seperate hosts. However, there is a lot more that we can do with this technique. An NN can be split any number of times without affecting the accuracy of the model.\n",
    "\n",
    "<b>Description: </b>Here we define a class which can procees a SplitNN of any number of layers. All it needs is a list of distributed models and their optimizers. \n",
    "\n",
    "\n",
    "\n",
    "<img src=\"images/MultiSplitNN.png\" width=\"20%\">\n",
    "\n",
    "\n",
    "In this tutorial, we demonstrate the SplitNN class with a 3 segment distribution [[1](https://arxiv.org/abs/1812.00564)]. This time;\n",
    "\n",
    "- <b>Alice</b>\n",
    "    - Has Model Segment 1\n",
    "    - Has the handwritten images\n",
    "- <b>Bob</b>\n",
    "    - Has model Segment 2\n",
    "- <b>Claire</b> \n",
    "    - Has Model Segment 3\n",
    "    - Has the image labels\n",
    "    \n",
    "We use the exact same model as we used in the previous tutorial, only this time we are splitting over 3 hosts, not two. However, we see the same loss being reported as there is <b>no reduction in accuracy</b> when training in this way. While we only use 3 models this can be done for any arbitrary number of models.\n",
    "\n",
    "Author:\n",
    "- Adam J Hall - Twitter: [@AJH4LL](https://twitter.com/AJH4LL) · GitHub:  [@H4LL](https://github.com/H4LL)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SplitNN:\n",
    "    def __init__(self, models, optimizers):\n",
    "        self.models = models\n",
    "        self.optimizers = optimizers\n",
    "        \n",
    "    def forward(self, x):\n",
    "        a = []\n",
    "        remote_a = []\n",
    "        \n",
    "        a.append(models[0](x))\n",
    "        if a[-1].location == models[1].location:\n",
    "            remote_a.append(a[-1].detach().requires_grad_())\n",
    "        else:\n",
    "            remote_a.append(a[-1].detach().move(models[1].location).requires_grad_())\n",
    "\n",
    "        i=1    \n",
    "        while i < (len(models)-1):\n",
    "            \n",
    "            a.append(models[i](remote_a[-1]))\n",
    "            if a[-1].location == models[i+1].location:\n",
    "                remote_a.append(a[-1].detach().requires_grad_())\n",
    "            else:\n",
    "                remote_a.append(a[-1].detach().move(models[i+1].location).requires_grad_())\n",
    "            \n",
    "            i+=1\n",
    "        \n",
    "        a.append(models[i](remote_a[-1]))\n",
    "        self.a = a\n",
    "        self.remote_a = remote_a\n",
    "        \n",
    "        return a[-1]\n",
    "    \n",
    "    def backward(self):\n",
    "        a=self.a\n",
    "        remote_a=self.remote_a\n",
    "        optimizers = self.optimizers\n",
    "        \n",
    "        i= len(models)-2   \n",
    "        while i > -1:\n",
    "            if remote_a[i].location == a[i].location:\n",
    "                grad_a = remote_a[i].grad.copy()\n",
    "            else:\n",
    "                grad_a = remote_a[i].grad.copy().move(a[i].location)\n",
    "            a[i].backward(grad_a)\n",
    "            i-=1\n",
    "\n",
    "    \n",
    "    def zero_grads(self):\n",
    "        for opt in optimizers:\n",
    "            opt.zero_grad()\n",
    "        \n",
    "    def step(self):\n",
    "        for opt in optimizers:\n",
    "            opt.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "from time import time\n",
    "from torchvision import datasets, transforms\n",
    "from torch import nn, optim\n",
    "import syft as sy\n",
    "import time\n",
    "hook = sy.TorchHook(torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                              transforms.Normalize((0.5,), (0.5,)),\n",
    "                              ])\n",
    "trainset = datasets.MNIST('mnist', download=True, train=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "# Define our model segments\n",
    "\n",
    "input_size = 784\n",
    "hidden_sizes = [128, 640]\n",
    "output_size = 10\n",
    "\n",
    "models = [\n",
    "    nn.Sequential(\n",
    "                nn.Linear(input_size, hidden_sizes[0]),\n",
    "                nn.ReLU(),\n",
    "    ),\n",
    "    nn.Sequential(\n",
    "                nn.Linear(hidden_sizes[0], hidden_sizes[1]),\n",
    "                nn.ReLU(),\n",
    "    ),\n",
    "    nn.Sequential(\n",
    "                nn.Linear(hidden_sizes[1], output_size),\n",
    "                nn.LogSoftmax(dim=1)\n",
    "    )\n",
    "]\n",
    "\n",
    "# Create optimisers for each segment and link to them\n",
    "optimizers = [\n",
    "    optim.SGD(model.parameters(), lr=0.03,)\n",
    "    for model in models\n",
    "]\n",
    "\n",
    "# create some workers\n",
    "alice = sy.VirtualWorker(hook, id=\"alice\")\n",
    "bob = sy.VirtualWorker(hook, id=\"bob\")\n",
    "claire = sy.VirtualWorker(hook, id=\"claire\")\n",
    "\n",
    "# Send Model Segments to model locations\n",
    "model_locations = [alice, bob, claire]\n",
    "for model, location in zip(models, model_locations):\n",
    "    model.send(location)\n",
    "\n",
    "#Instantiate a SpliNN class with our distributed segments and their respective optimizers\n",
    "splitNN =  SplitNN(models, optimizers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(x, target, splitNN):\n",
    "    \n",
    "    #1) Zero our grads\n",
    "    splitNN.zero_grads()\n",
    "    \n",
    "    #2) Make a prediction\n",
    "    pred = splitNN.forward(x)\n",
    "    \n",
    "    #3) Figure out how much we missed by\n",
    "    criterion = nn.NLLLoss()\n",
    "    loss = criterion(pred, target)\n",
    "    \n",
    "    #4) Backprop the loss on the end layer\n",
    "    loss.backward()\n",
    "    \n",
    "    #5) Feed Gradients backward through the nework\n",
    "    splitNN.backward()\n",
    "    \n",
    "    #6) Change the weights\n",
    "    splitNN.step()\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(epochs):\n",
    "    running_loss = 0\n",
    "    for images, labels in trainloader:\n",
    "        images = images.send(models[0].location)\n",
    "        images = images.view(images.shape[0], -1)\n",
    "        labels = labels.send(models[-1].location)\n",
    "        loss = train(images, labels, splitNN)\n",
    "        running_loss += loss.get()\n",
    "\n",
    "    else:\n",
    "        print(\"Epoch {} - Training loss: {}\".format(i, running_loss/len(trainloader)))"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
